import os
import shutil
import zipfile
from pathlib import Path
from tempfile import TemporaryDirectory

def replace_images_in_docx(docx_path, new_images_dir, output_path):
    """
    替换 Word 文档中的图片
    :param docx_path: 原 Word 文档路径
    :param new_images_dir: 新图片目录（需包含同名文件）
    :param output_path: 输出文档路径
    """
    # 创建临时工作目录
    with TemporaryDirectory() as tmp_dir:
        tmp_dir = Path(tmp_dir)
        
        # 1. 解压原文档到临时目录
        with zipfile.ZipFile(docx_path, 'r') as zip_ref:
            zip_ref.extractall(tmp_dir)
        
        # 2. 定位图片文件夹
        media_dir = tmp_dir / "word" / "media"
        if not media_dir.exists():
            raise FileNotFoundError("未找到图片目录 word/media/")
        
        # 3. 替换图片
        replaced = []
        for img_file in Path(new_images_dir).glob("*"):
            if img_file.suffix.lower() in ('.png', '.jpg', '.jpeg', '.gif'):
                target_path = media_dir / img_file.name
                if target_path.exists():
                    shutil.copy(img_file, target_path)
                    replaced.append(img_file.name)
        
        if not replaced:
            print("警告：未找到可替换的图片（需保持文件名一致）")
            return

        # 4. 重新打包成新文档
        with zipfile.ZipFile(output_path, 'w', zipfile.ZIP_DEFLATED) as new_zip:
            for root, dirs, files in os.walk(tmp_dir):
                for file in files:
                    file_path = Path(root) / file
                    arcname = file_path.relative_to(tmp_dir)
                    new_zip.write(file_path, arcname)
        
        print(f"成功替换图片: {', '.join(replaced)}")
        print(f"新文档已生成: {output_path}")

if __name__ == "__main__":
    # 使用示例
    original_docx = "www.docx"
    new_images_folder = "1"
    output_docx = "替换后的文档.docx"

    replace_images_in_docx(original_docx, new_images_folder, output_docx)